import os
import requests
from volcenginesdkarkruntime import Ark

# 文本向量API的URL
VOLC_EMBEDDINGS_API_URL = os.getenv("VOLC_EMBEDDINGS_API_URL")
# 视觉向量API的URL
VOLC_EMBEDDINGS_VISION_API_URL = os.getenv("VOLC_EMBEDDINGS_VISION_API_URL")
# API密钥
VOLC_API_KEY = os.getenv("VOLC_API_KEY")

def get_doc_embedding(doc_content):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {VOLC_API_KEY}",
    }
    payload = {
        "model": "doubao-embedding-text-240715",
        "input": doc_content,
    }
    response = requests.post(
        VOLC_EMBEDDINGS_API_URL, json=payload, headers=headers
    )
    if response.status_code == 200:
        # 解析响应的json数据
        data = response.json()
        # 提取嵌入向量
        embedding = data["data"][0]["embedding"]
        return embedding
    else:
        raise Exception(f"Embedding API error{response.text}")

# 定义获取图片嵌入向量的函数
def get_image_embedding(image_url):
    client = Ark(api_key=VOLC_API_KEY)

    response = client.multimodal_embeddings.create(
        model="doubao-embedding-vision-250615",
        encoding_format="float",
        input=[
            {
                "image_url": {
                    "url": image_url,
                },
                "type": "image_url"
            }
        ]
    )
    # 解析响应的json数据
    data = response.model_dump()
    # 提取嵌入向量
    embedding = data["data"]["embedding"]
    return embedding

